[DeepSeek-V4.1] Optimize DSpark verify and MoE kernels on Blackwell - #38879

"""Fixed-input, real-acceptance BS1 measurements on the owned server only."""
import argparse
import os
import hashlib
import json
import re
import time
from pathlib import Path

import requests
from tokenizers import Tokenizer


def save(path, value):
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--out', type=Path, required=True)
    parser.add_argument('--server-log', type=Path, required=True)
    parser.add_argument('--url', default='http://127.0.0.1:30021')
    parser.add_argument('--repeat', type=int, default=3)
    parser.add_argument('--profile', action='store_true')
    args = parser.parse_args()
    tok = Tokenizer.from_file(os.path.join(os.environ['MODEL_PATH'], 'tokenizer.json'))
    prefix = tok.encode('Read these notes and summarize them.\n', add_special_tokens=False).ids
    unit = tok.encode('The observatory records the temperature, wind, and rainfall each day. Researchers compare the measurements across seasons.\n', add_special_tokens=False).ids
    suffix = tok.encode('\nWrite a detailed summary of the notes above.\n', add_special_tokens=False).ids

    def post(route, body=None, timeout=1800):
        r = requests.post(args.url+route, json=body or {}, timeout=timeout)
        r.raise_for_status()
        return r

    def warm(ids, rate):
        # The final SSE response can precede scheduler request cleanup.
        # Use the server's bounded idle wait, outside the measured interval.
        post('/flush_cache?timeout=30', timeout=60)
        if rate:
            r = post('/generate', {'input_ids':ids[:int(len(ids)*rate)], 'sampling_params':{'temperature':0, 'max_new_tokens':1, 'ignore_eos':True}})
            assert r.json()['meta_info']['completion_tokens']==1

    results=[]
    for length, out_len, rate in ((4096,1024,0.),):
        n = length-len(prefix)-len(suffix)
        ids = prefix + (unit*((n+len(unit)-1)//len(unit)))[:n] + suffix
        assert len(ids)==length
        save(args.out/f'input-{length}.json',ids)
        input_hash=hashlib.sha256(json.dumps(ids).encode()).hexdigest()
        for rep in range(args.repeat+1):
            warm(ids,rate)
            log_offset=args.server_log.stat().st_size
            body={'input_ids':ids,'sampling_params':{'temperature':0,'max_new_tokens':out_len,'ignore_eos':True,'stream_interval':1},'stream':True}
            start=time.perf_counter();first=None;first_count=None;last=None
            with requests.post(args.url+'/generate',json=body,stream=True,timeout=(30,1800)) as r:
                r.raise_for_status()
                for line in r.iter_lines():
                    if not line.startswith(b'data: '):continue
                    raw=line[6:]
                    if raw==b'[DONE]':break
                    last=json.loads(raw)
                    if 'error' in last:raise RuntimeError(last)
                    count=last.get('meta_info',{}).get('completion_tokens',0)
                    if first is None and count>0:first=time.perf_counter();first_count=count
            end=time.perf_counter()
            assert last is not None and first is not None,last
            meta=last['meta_info']
            assert meta['completion_tokens']==out_len and meta['prompt_tokens']==length,meta
            # Required evidence of actual speculative execution and acceptance.
            assert meta.get('spec_verify_ct',0)>0 and meta.get('spec_accept_length',0)>0,meta
            with args.server_log.open() as f:f.seek(log_offset);log=f.read()
            samples=[float(v) for v in re.findall(r'gen throughput \(token/s\):\s*([0-9.]+)',log)]
            result=dict(input_tokens=length,output_tokens=out_len,requested_cache_hit_rate=rate,repeat=rep,warmup=(rep==0),input_sha256=input_hash,elapsed_s=end-start,ttft_s=first-start,first_stream_completion_tokens=first_count,post_first_event_tps=(out_len-first_count)/(end-first),full_request_tps=out_len/(end-start),decode_log_samples=samples,response=last)
            results.append(result);save(args.out/'performance.json',results)
            print('PERF',length,rep,result['post_first_event_tps'],'accept',meta['spec_accept_length'],'cached',meta.get('cached_tokens'),flush=True)

        if args.profile:
            for activities in (['GPU'],):
                if length==131072 and activities==['CPU','GPU']:continue
                warm(ids,rate)
                tag=f'{length}-'+('-'.join(activities)).lower()
                folder=args.out/('profile-'+tag);folder.mkdir()
                request=dict(output_dir=str(folder),num_steps=20,activities=activities,profile_by_stage=True,profile_stages=['decode'],with_stack=(activities==['CPU','GPU']),record_shapes=(activities==['CPU','GPU']),profile_id=args.out.name+'-'+tag)
                save(folder/'request.json',request)
                post('/start_profile',request,timeout=60)
                response=post('/generate',{'input_ids':ids,'sampling_params':{'temperature':0,'max_new_tokens':200,'ignore_eos':True}}).json()
                save(folder/'response.json',response)
                # At most six tokens per verify step: 200 output tokens
                # always cover the 20 requested decode steps. Export is automatic.
                traces=list(folder.rglob('*.trace.json*'))
                assert len(traces)>=4,(tag,'Missing TP traces',traces)
                save(folder/'trace-files.json',[str(p) for p in traces])
                print('PROFILE',tag,len(traces),flush=True)


if __name__=='__main__':
    main()

"""BS64 serving regression measurement; raw requests and log intervals retained."""
import argparse
import os
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
from pathlib import Path
import re
import statistics
import threading
import time
import urllib.request

from tokenizers import Tokenizer


def save(path, value):
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + '\n')


def post(url, route, body=None):
    request = urllib.request.Request(url + route, data=json.dumps(body or {}).encode(),
                                     headers={'Content-Type': 'application/json'})
    with urllib.request.urlopen(request, timeout=1800) as response:
        if route.startswith(('/flush_cache', '/start_profile', '/stop_profile')):
            return response.read().decode()
        return json.load(response)


def decode_samples(log):
    rows = []
    previous_bs = None
    for line in log.replace('\r', '\n').splitlines():
        if 'Prefill batch' in line:
            previous_bs = None
        if 'Decode batch' not in line:
            continue
        bs = re.search(r'#running-req: (\d+)', line)
        tps = re.search(r'gen throughput \(token/s\): ([0-9.]+)', line)
        al = re.search(r'accept len: ([0-9.]+)', line)
        graph = re.search(r'(?:CUDA graph|cuda graph|cuda_graph|full graph): (True|False)', line)
        if not bs or not tps:
            continue
        batch = int(bs.group(1))
        rows.append(dict(batch_size=batch, tps=float(tps.group(1)),
                         accept_length=float(al.group(1)) if al else None,
                         graph=graph.group(1) == 'True' if graph else None,
                         eligible=(batch == 64 and previous_bs == 64), line=line))
        previous_bs = batch
    return rows


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--out', type=Path, required=True)
    parser.add_argument('--server-log', type=Path, required=True)
    parser.add_argument('--url', default='http://127.0.0.1:30021')
    parser.add_argument('--no-spec', action='store_true')
    args = parser.parse_args()
    tok = Tokenizer.from_file(os.path.join(os.environ['MODEL_PATH'], 'tokenizer.json'))
    unit = tok.encode('The observatory records the temperature, wind, and rainfall each day. Researchers compare the measurements across seasons.\n', add_special_tokens=False).ids
    suffix = tok.encode('\nWrite a detailed summary of the notes above.\n', add_special_tokens=False).ids
    inputs = []
    for i in range(64):
        prefix = tok.encode(f'Notebook {i:03d}. Read these notes and summarize them.\n', add_special_tokens=False).ids
        n = 4096 - len(prefix) - len(suffix)
        inputs.append(prefix + (unit * ((n + len(unit) - 1) // len(unit)))[:n] + suffix)
    assert len({tuple(ids) for ids in inputs}) == 64
    assert all(len(ids) == 4096 for ids in inputs)
    save(args.out/'inputs.json', inputs)
    input_hash = hashlib.sha256(json.dumps(inputs).encode()).hexdigest()
    waves = []
    for rep in range(4):
        warmup = rep == 0
        output_length = 256 if warmup else 2048
        post(args.url, '/flush_cache?timeout=30')
        offset = args.server_log.stat().st_size
        barrier = threading.Barrier(65)

        def request_one(i):
            body = {'input_ids': inputs[i], 'sampling_params': {
                'temperature': 0, 'max_new_tokens': output_length, 'ignore_eos': True}}
            barrier.wait()
            begin = time.perf_counter()
            response = post(args.url, '/generate', body)
            elapsed = time.perf_counter() - begin
            meta = response['meta_info']
            assert meta['completion_tokens'] == output_length, meta
            assert meta['prompt_tokens'] == 4096, meta
            if not args.no_spec:
                assert meta.get('spec_verify_ct', 0) > 0, meta
                assert meta.get('spec_accept_length', 0) > 0, meta
            return dict(index=i, elapsed_s=elapsed, response=response)

        with ThreadPoolExecutor(max_workers=64) as executor:
            futures = [executor.submit(request_one, i) for i in range(64)]
            start = time.perf_counter()
            barrier.wait()
            responses = [future.result() for future in futures]
            elapsed = time.perf_counter() - start
        save(args.out/f'wave-{rep}-responses.json', responses)
        # Drain scheduler request cleanup before the next cache flush.
        time.sleep(1)
        with args.server_log.open('rb') as log:
            log.seek(offset)
            text = log.read().decode(errors='replace')
        (args.out/f'wave-{rep}-server.log').write_text(text)
        rows = decode_samples(text)
        eligible = [row for row in rows if row['eligible']]
        if not warmup:
            assert len(eligible) >= 3, ('Insufficient actual BS64 intervals', rows)
            assert all(row['graph'] is True for row in eligible), eligible[:3]
        completed = sum(r['response']['meta_info']['completion_tokens'] for r in responses)
        verifies = sum(r['response']['meta_info'].get('spec_verify_ct', 0) for r in responses)
        wave = dict(repeat=rep, warmup=warmup, requests=64, input_tokens=4096,
                    output_tokens_per_request=output_length, input_sha256=input_hash,
                    elapsed_s=elapsed, aggregate_full_request_tps=completed/elapsed,
                    weighted_accept_length=completed/verifies if verifies else None,
                    bs64_decode_median_tps=statistics.median(row['tps'] for row in eligible) if eligible else None,
                    bs64_decode_intervals=len(eligible), decode_intervals=rows)
        waves.append(wave)
        save(args.out/'performance.json', waves)
        print('WAVE', rep, 'warmup', warmup, 'actual BS64 decode', wave['bs64_decode_median_tps'],
              'intervals', len(eligible), 'full-request', wave['aggregate_full_request_tps'],
              'AL', wave['weighted_accept_length'], flush=True)


if __name__ == '__main__':
    main()

"""Pinned legacy five-shot GSM8K prompts/scorer, with complete response/error records."""
import argparse,hashlib,json,time
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor,as_completed
import requests
from sglang.test.simple_eval_mixed_prefix_gsm8k import get_few_shot_examples,get_one_example,get_answer_value

def main():
 p=argparse.ArgumentParser();p.add_argument('--out',type=Path,required=True);p.add_argument('--threads',type=int,required=True);p.add_argument('--count',type=int,default=1314);p.add_argument('--url',default='http://127.0.0.1:30021');a=p.parse_args()
 a.out.mkdir(exist_ok=True);data=Path(__file__).with_name('gsm8k-test.jsonl');rows=[json.loads(l) for l in data.read_text().splitlines()];assert len(rows)==1319
 prefix=get_few_shot_examples(rows,5);ids=list(range(5,min(1319,5+a.count)))
 models=requests.get(a.url+'/v1/models',timeout=30);models.raise_for_status();model=models.json()['data'][0]['id']
 def run(i):
  prompt=prefix+get_one_example(rows,i,False);start=time.perf_counter()
  rec=dict(index=i,prompt=prompt,prompt_sha256=hashlib.sha256(prompt.encode()).hexdigest(),expected=get_answer_value(rows[i]['answer']))
  try:
   r=requests.post(a.url+'/v1/chat/completions',json=dict(model=model,messages=[dict(role='user',content=prompt)],temperature=0,top_p=1,max_tokens=4096,seed=0,return_meta_info=True),timeout=(30,1800));r.raise_for_status();j=r.json();rec['response']=j
   ch=j['choices'][0];content=ch['message'].get('content') or '';rec.update(answer=get_answer_value(content),correct=get_answer_value(content)==rec['expected'],truncated=ch['finish_reason']=='length',empty=not bool(content.strip()))
  except Exception as e:rec.update(error=repr(e),correct=False,truncated=False,empty=True)
  rec['elapsed_s']=time.perf_counter()-start;return rec
 results=[];begin=time.perf_counter()
 with (a.out/'samples.jsonl').open('w') as f,ThreadPoolExecutor(max_workers=a.threads) as pool:
  futures=[pool.submit(run,i) for i in ids]
  for future in as_completed(futures):
   r=future.result();results.append(r);f.write(json.dumps(r,ensure_ascii=False)+'\n');f.flush()
   if len(results)%25==0 or len(results)==len(ids):print('GSM',len(results),'/',len(ids),'correct',sum(x['correct'] for x in results),'errors',sum('error' in x for x in results),flush=True)
 summary=dict(count=len(results),correct=sum(r['correct'] for r in results),score=sum(r['correct'] for r in results)/len(results),errors=sum('error' in r for r in results),truncated=sum(r['truncated'] for r in results),empty=sum(r['empty'] for r in results),threads=a.threads,elapsed_s=time.perf_counter()-begin,dataset_sha256=hashlib.sha256(data.read_bytes()).hexdigest(),num_shots=5,held_out_indices=ids,temperature=0,top_p=1,seed=0,max_tokens=4096)
 (a.out/'summary.json').write_text(json.dumps(summary,indent=2));print(json.dumps(summary),flush=True)
 assert summary['errors']==0,summary
 assert summary['score']>=.9,summary
if __name__=='__main__':main()